import numpy as np
import pandas as pd
from utils import prepare_Y
from sklearn.linear_model import LinearRegression


ydata = prepare_Y().T
ldata = pd.read_excel("data/Ldata/劳动力数量.xlsx",index_col=0)
kdata = pd.read_excel("data/k.xlsx",index_col=0)
print(ldata)

df = pd.concat([ydata['制造'], kdata['制造'], ldata['制造']], axis=1)
df.columns = ['Y', 'K', 'L']
# 假设您的数据已存储在DataFrame中，列名为'Y', 'K', 'L'
df.sort_index(inplace=True)
# 1. 对数据取自然对数

df = df[df.index >= 2004] # 部分2003之前的行业有缺失值Y
df = df[df.index <= 2023] # 部分2024年的数据有缺失值

# 使用规模报酬不变约束 (alpha + beta = 1)
# 通过重新参数化：ln(Y/L) = ln(A) + alpha*ln(K/L)
df['ln_Y_L'] = np.log(df['Y'] / df['L'])
df['ln_K_L'] = np.log(df['K'] / df['L'])

# 准备数据
X = df[['ln_K_L']]  # 自变量现在只有一个：ln(K/L)
y = df['ln_Y_L']    # 因变量是ln(Y/L)

# 执行线性回归
model = LinearRegression()
model.fit(X, y)

# 获取估计的系数
alpha = model.coef_[0]   # 资本产出弹性
beta = 1 - alpha         # 劳动产出弹性（根据约束alpha+beta=1）
ln_A = model.intercept_  # ln(A)
A = np.exp(ln_A)         # 全要素生产率

print(f"资本产出弹性 (α): {alpha:.4f}")
print(f"劳动产出弹性 (β): {beta:.4f}")
print(f"α + β = {alpha + beta:.4f}")  # 验证约束是否生效
print(f"全要素生产率 (A): {A:.4f}")